"""
预处理的GMRES求解器

该模块实现了用于求解线性方程组的预处理GMRES方法：
1. 基于物理的预处理器
2. 多重网格预处理器
3. 自适应预处理器选择
4. 改进的收敛策略
"""

import numpy as np
from typing import Tuple, Optional, Callable, List, Dict
from scipy.sparse import spmatrix, csr_matrix, diags
from scipy.sparse.linalg import spsolve
import logging

logger = logging.getLogger(__name__)

class BasePreconditioner:
    """预处理器基类"""
    
    def build_preconditioner(self, A: spmatrix) -> spmatrix:
        """构建预处理器矩阵"""
        raise NotImplementedError
        
    def apply(self, x: np.ndarray) -> np.ndarray:
        """应用预处理器"""
        raise NotImplementedError

class PhysicalPreconditioner(BasePreconditioner):
    """基于物理的预处理器"""
    
    def __init__(self, n_axial: int, n_channels: int, include_lateral: bool = False):
        """
        初始化预处理器
        
        参数:
            n_axial: 轴向节点数
            n_channels: 子通道数量
            include_lateral: 是否包含横向流动
        """
        self.n_axial = n_axial
        self.n_channels = n_channels
        self.include_lateral = include_lateral
        self.M = None
        
    def build_preconditioner(self, A: spmatrix) -> spmatrix:
        """
        构建预处理器矩阵
        
        参数:
            A: 原始系数矩阵
            
        返回:
            预处理器矩阵
        """
        # 将稀疏矩阵转换为CSR格式
        A_csr = csr_matrix(A)
        
        # 获取矩阵的行、列索引和数据
        rows = A_csr.indptr
        cols = A_csr.indices
        data = A_csr.data
        
        # 创建预处理器矩阵
        M_data = []
        M_rows = []
        M_cols = []
        
        for i in range(self.n_axial * self.n_channels):
            row_start = rows[i]
            row_end = rows[i+1]
            
            row_cols = cols[row_start:row_end]
            row_data = data[row_start:row_end]
            
            axial_pos = i // self.n_channels
            channel = i % self.n_channels
            
            for j, col in enumerate(row_cols):
                col_axial = col // self.n_channels
                col_channel = col % self.n_channels
                
                # 保留中心单元和轴向相邻单元的元素
                if (col_channel == channel and abs(col_axial - axial_pos) <= 1):
                    M_data.append(row_data[j])
                    M_rows.append(i)
                    M_cols.append(col)
                # 如果包含横向流动，保留横向相邻单元的元素
                elif (self.include_lateral and 
                      col_axial == axial_pos and 
                      abs(col_channel - channel) == 1):
                    M_data.append(row_data[j] * 0.5)  # 横向元素权重减半
                    M_rows.append(i)
                    M_cols.append(col)
                    
        self.M = csr_matrix((M_data, (M_rows, M_cols)), shape=A_csr.shape)
        return self.M
        
    def apply(self, x: np.ndarray) -> np.ndarray:
        """应用预处理器"""
        if self.M is None:
            raise RuntimeError("预处理器矩阵尚未构建")
        return spsolve(self.M, x)

class MultiGridPreconditioner(BasePreconditioner):
    """多重网格预处理器"""
    
    def __init__(self, n_axial: int, n_channels: int, levels: int = 3):
        """
        初始化多重网格预处理器
        
        参数:
            n_axial: 轴向节点数
            n_channels: 子通道数量
            levels: 网格层数
        """
        self.n_axial = n_axial
        self.n_channels = n_channels
        self.levels = levels
        self.operators = []  # 存储各层次的算子
        self.M = None
        
    def build_preconditioner(self, A: spmatrix) -> spmatrix:
        """构建多重网格预处理器"""
        self.operators = []
        current_A = csr_matrix(A)
        
        # 构建各层次的算子
        for level in range(self.levels):
            # 构建限制和延拓算子
            R = self._build_restriction(current_A.shape[0])
            P = R.T
            
            # 计算粗网格算子
            coarse_A = R @ current_A @ P
            
            self.operators.append({
                'A': current_A,
                'R': R,
                'P': P
            })
            
            current_A = coarse_A
            
        # 最粗网格直接求解
        self.operators.append({'A': current_A})
        
        # 存储原始矩阵作为预处理器矩阵
        self.M = A
        return self.M
        
    def apply(self, x: np.ndarray) -> np.ndarray:
        """应用多重网格预处理"""
        return self._v_cycle(x, 0)
        
    def _v_cycle(self, b: np.ndarray, level: int) -> np.ndarray:
        """执行V-循环"""
        if level == self.levels:
            # 最粗网格直接求解
            return spsolve(self.operators[level]['A'], b)
            
        # 前光滑
        x = self._smooth(self.operators[level]['A'], b)
        
        # 计算残差
        r = b - self.operators[level]['A'] @ x
        
        # 限制到粗网格
        r_coarse = self.operators[level]['R'] @ r
        
        # 递归求解粗网格
        e_coarse = self._v_cycle(r_coarse, level + 1)
        
        # 延拓回细网格
        x = x + self.operators[level]['P'] @ e_coarse
        
        # 后光滑
        x = self._smooth(self.operators[level]['A'], b, x)
        
        return x
        
    def _build_restriction(self, n: int) -> spmatrix:
        """构建限制算子"""
        # 简单的注入限制算子
        n_coarse = n // 2
        data = np.ones(n_coarse)
        indices = np.arange(0, n, 2)
        indptr = np.arange(n_coarse + 1)
        return csr_matrix((data, indices, indptr), shape=(n_coarse, n))
        
    def _smooth(self, A: spmatrix, b: np.ndarray, 
               x0: Optional[np.ndarray] = None, 
               n_iter: int = 2) -> np.ndarray:
        """Gauss-Seidel光滑"""
        if x0 is None:
            x0 = np.zeros_like(b)
        x = x0.copy()
        
        for _ in range(n_iter):
            for i in range(A.shape[0]):
                r = b[i] - A[i,:].dot(x)
                x[i] = x[i] + r / A[i,i]
                
        return x

class AdaptivePreconditioner(BasePreconditioner):
    """自适应预处理器"""
    
    def __init__(self, n_axial: int, n_channels: int):
        """初始化自适应预处理器"""
        self.n_axial = n_axial
        self.n_channels = n_channels
        self.preconditioners = {
            'physical': PhysicalPreconditioner(n_axial, n_channels),
            'physical_lateral': PhysicalPreconditioner(n_axial, n_channels, True),
            'multigrid': MultiGridPreconditioner(n_axial, n_channels)
        }
        self.current = None
        self.M = None
        
    def build_preconditioner(self, A: spmatrix) -> spmatrix:
        """根据矩阵特征选择合适的预处理器"""
        # 估计条件数
        diag = np.abs(A.diagonal())
        condition_estimate = np.max(diag) / np.min(diag)
        
        # 分析矩阵结构
        n_nonzero = A.getnnz()
        density = n_nonzero / (A.shape[0] * A.shape[1])
        
        # 选择预处理器
        if condition_estimate > 1e4:
            logger.info("选择多重网格预处理器")
            self.current = 'multigrid'
        elif density > 0.1:
            logger.info("选择包含横向流动的物理预处理器")
            self.current = 'physical_lateral'
        else:
            logger.info("选择基本物理预处理器")
            self.current = 'physical'
            
        self.M = self.preconditioners[self.current].build_preconditioner(A)
        return self.M
        
    def apply(self, x: np.ndarray) -> np.ndarray:
        """应用当前选择的预处理器"""
        if self.current is None:
            raise RuntimeError("尚未选择预处理器")
        return self.preconditioners[self.current].apply(x)

class PGMRESSolver:
    """预处理的GMRES求解器"""
    
    def __init__(self, preconditioner: BasePreconditioner):
        """
        初始化求解器
        
        参数:
            preconditioner: 预处理器对象
        """
        self.preconditioner = preconditioner
        
    def solve(self, A: spmatrix, b: np.ndarray, x0: Optional[np.ndarray] = None,
             max_iter: int = 100, tol: float = 1e-6, 
             restart: int = 30) -> Tuple[np.ndarray, bool, int]:
        """
        求解线性方程组 Ax = b
        
        参数:
            A: 系数矩阵
            b: 右端向量
            x0: 初始猜测解
            max_iter: 最大迭代次数
            tol: 收敛容差
            restart: 重启步数
            
        返回:
            (解向量, 是否收敛, 迭代次数)
        """
        # 构建预处理器
        M = self.preconditioner.build_preconditioner(A)
        
        # 初始化
        n = A.shape[0]
        if x0 is None:
            x0 = np.zeros(n)
            
        total_iter = 0
        x = x0.copy()
        
        while total_iter < max_iter:
            # 计算初始残差
            r0 = b - A @ x
            beta = np.linalg.norm(r0)
            
            if beta < tol:
                logger.info(f"PGMRES在{total_iter}次迭代后收敛")
                return x, True, total_iter
                
            # 初始化Krylov子空间基向量
            V = np.zeros((n, restart + 1))
            H = np.zeros((restart + 1, restart))
            V[:,0] = r0 / beta
            
            # 旋转矩阵的sine和cosine值
            cs = np.zeros(restart)
            sn = np.zeros(restart)
            
            # 残差范数向量
            e1 = np.zeros(restart + 1)
            e1[0] = beta
            
            # GMRES迭代
            for j in range(restart):
                # Arnoldi过程
                V[:,j+1], H[:j+2,j] = self._arnoldi(A, M, V[:,:j+1], j)
                
                # 应用Givens旋转
                for i in range(j):
                    temp = cs[i] * H[i,j] + sn[i] * H[i+1,j]
                    H[i+1,j] = -sn[i] * H[i,j] + cs[i] * H[i+1,j]
                    H[i,j] = temp
                    
                # 计算新的Givens旋转
                cs[j], sn[j] = self._givens_rotation(H[j,j], H[j+1,j])
                
                # 应用新的旋转
                H[j,j] = cs[j] * H[j,j] + sn[j] * H[j+1,j]
                H[j+1,j] = 0
                
                # 更新残差范数
                e1[j+1] = -sn[j] * e1[j]
                e1[j] = cs[j] * e1[j]
                
                # 检查收敛性
                residual = abs(e1[j+1])
                total_iter += 1
                
                if residual < tol * beta:
                    # 计算解向量
                    y = np.linalg.solve(H[:j+1,:j+1], e1[:j+1])
                    x = x + V[:,:j+1] @ y
                    logger.info(f"PGMRES在{total_iter}次迭代后收敛")
                    return x, True, total_iter
                    
                if total_iter >= max_iter:
                    break
                    
            # 重启：使用当前最佳解作为新的初始值
            y = np.linalg.solve(H[:restart,:restart], e1[:restart])
            x = x + V[:,:restart] @ y
            
        logger.warning(f"PGMRES在达到最大迭代次数{max_iter}后未收敛")
        return x, False, total_iter
        
    def _arnoldi(self, A: spmatrix, M: spmatrix, V: np.ndarray, 
                j: int) -> Tuple[np.ndarray, np.ndarray]:
        """
        执行Arnoldi过程的一步
        
        参数:
            A: 系数矩阵
            M: 预处理器矩阵
            V: Krylov子空间基向量
            j: 当前迭代步
            
        返回:
            (新的基向量, Hessenberg矩阵的新列)
        """
        # 应用预处理的矩阵
        w = self.preconditioner.apply(A @ V[:,j])
        
        # Modified Gram-Schmidt正交化
        h = np.zeros(j + 2)
        for i in range(j + 1):
            h[i] = np.dot(w, V[:,i])
            w = w - h[i] * V[:,i]
            
        h[j+1] = np.linalg.norm(w)
        if h[j+1] != 0:
            w = w / h[j+1]
            
        return w, h
        
    def _givens_rotation(self, h1: float, h2: float) -> Tuple[float, float]:
        """
        计算Givens旋转矩阵的参数
        
        参数:
            h1: Hessenberg矩阵的对角元素
            h2: Hessenberg矩阵的次对角元素
            
        返回:
            (cosine, sine)
        """
        if h2 == 0:
            return 1, 0
            
        if abs(h2) > abs(h1):
            tau = -h1 / h2
            sn = 1 / np.sqrt(1 + tau**2)
            cs = tau * sn
        else:
            tau = -h2 / h1
            cs = 1 / np.sqrt(1 + tau**2)
            sn = tau * cs
            
        return cs, sn